package com.farm2.llm.client.impl;

import com.farm2.llm.client.inter.Farm2LlmClientInter;
import com.farm2.llm.client.handle.inter.Farm2LlmMessageHandleInter;
import com.farm2.llm.client.msg.ChatMsg;
import com.plexpt.chatgpt.ChatGPTStream;
import com.plexpt.chatgpt.entity.chat.ChatCompletion;
import com.plexpt.chatgpt.entity.chat.Message;
import com.plexpt.chatgpt.listener.AbstractStreamListener;
import okhttp3.sse.EventSource;
import org.farm2.llm.domain.LlmClient;
import org.farm2.tools.json.FarmJsons;
import org.farm2.tools.web.dto.SendMsgDto;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

@Component
public class Farm2LlmPlexPtClientImpl implements Farm2LlmClientInter {
    private LlmClient paras = null;

    @Override
    public String getName() {
        return "PlexPt_ChatGPT_SDK";
    }

    @Override
    public void init(LlmClient client) {
        paras = client;
    }

    @Override
    public int getTokenSize() {
        return paras.getTokensize();
    }


    @Override
    public void sendMsg(ChatMsg msg, Farm2LlmMessageHandleInter handle) {
        if (paras == null) {
            throw new RuntimeException("当前模型客户端未初始化 paras=null");
        }
        ChatGPTStream chatGPTStream = ChatGPTStream.builder()
                .timeout(3000)
                .apiKey(paras.getApikey())
                .apiHost(paras.getBaseurl())
                .build()
                .init();

        List<Message> msgs = new ArrayList<>();
        for (SendMsgDto smsg : msg.getHisMsgs()) {
            msgs.add(Message.of(smsg.getMsg()));
            msgs.add(Message.ofAssistant(smsg.getSysMsg()));
        }
        msgs.add(Message.of(msg.getMsg()));

        System.out.println(FarmJsons.toFormatJson(msgs));

        chatGPTStream.streamChatCompletion(ChatCompletion.builder().model(paras.getModelkey())
                .messages(msgs)
                .build(), new AbstractStreamListener() {
            @Override
            public void onMsg(String s) {
                handle.onMsg(s);
            }

            @Override
            public void onClosed(EventSource eventSource) {
                handle.onClosed();
                eventSource.cancel();
            }

            @Override
            public void onError(Throwable throwable, String s) {
                handle.onError(throwable, s);
            }
        });
    }
}
